#!/usr/bin/env python
"""
Copyright Ian Ross Williams, 2012

    This file is part of ArcSecond.

    ArcSecond is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    ArcSecond is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with ArcSecond.  If not, see <http://www.gnu.org/licenses/>.
"""
import numpy as np
from numpy.fft import rfft2,irfft2
import Image,ImageOps
from scipy.optimize import fmin_l_bfgs_b,fmin_powell
import scipy.special as special
from scipy import factorial
from arcsecond import filters
import pywt
import os
import time

debugPath = None
#debugPath = "/home/ian/workspace/gtkapture/"
#debugPath = "/home/ian/Code/gtkapture/"
#debugPath = "/home/wi10008/Code/gtkapture/gtkapture/"

class MFBD:
    ( INITIAL_FIRSTFRAME,
      INITIAL_MEAN,
      INITIAL_MEDIAN,
      INITIAL_KALMAN) = range(4)
           
    #Shapelet PSF Multi-Frame Blind Deconvolution
    def MFBD(self,stack,shape,psfsize,blurRadius,
             stackExtra=None,callback=None,roi=None,
             tolerance=1./255.,
             EMiterations=10,Eiterations=10.,Miterations=5,
             methodInitialGuess=INITIAL_MEDIAN,
             shapeletOrder=3,underRelax=1.0, 
             lambdaTV=0, LRThreshold = 0, waveletFilter=False, 
             ignoreImagesWithStdDevBelow=1.,noOfFrames=1):

        if debugPath!=None:
            os.system("rm " + debugPath + "src/test/*")
        EMiterations=int(EMiterations)
        Eiterations=int(Eiterations)
        Miterations=int(Miterations)
        shapeletOrder=int(shapeletOrder)
        noOfFrames=int(noOfFrames)
        
        allFrames = stack.shape[0]
        mask = np.arange(stack.shape[0]) < noOfFrames
        sstack = stack[mask]
        if stackExtra!=None:
            for s in stackExtra: 
                s = s[mask]
            
        allPSF = np.zeros((noOfFrames,shape[0],shape[1]))
        tolerance = max(tolerance,1./(np.max(stack) - np.min(stack)))
        
        if roi==None: roi=(0,0,shape[1],shape[0])
        roi = (int(roi[0]),int(roi[1]),int(roi[2]),int(roi[3]))
        samplesize=min(roi[2]-roi[0],roi[3]-roi[1])
        sampleshape=(samplesize,samplesize)
        psfsize = min(psfsize,samplesize/2)
        psfshape = (psfsize,psfsize)
                                
        x0=int(max(roi[1],psfsize))
        x1=int(min(x0+samplesize,shape[0]-psfsize))
        y0=int(max(roi[0],psfsize))
        y1=int(min(y0+samplesize,shape[1]-psfsize))
        samplesize=min(x1-x0,y1-y0)

        px0,px1=x0-psfsize, x1+psfsize
        py0,py1=y0-psfsize, y1+psfsize

        psfmx, psfmy = (x1+x0)/2,(y1+y0)/2
        psfx0 = psfmx - psfsize/2
        psfy0 = psfmy - psfsize/2
        psfx1 = psfx0 + psfsize
        psfy1 = psfy0 + psfsize
        
        #Form an initial guess at the answer
        xn = sstack[0]
        xn = xn.reshape(shape)
        imX = Image.fromarray(xn.astype(np.int8),"L")
        
        if methodInitialGuess==self.INITIAL_MEAN:
            alive=callback(imX, "Forming initial guess using Mean stack" )
            if not alive: return None
            xn = np.mean(stack,axis=0)
            xn = xn.reshape(shape)
        if methodInitialGuess==self.INITIAL_MEDIAN:
            alive=callback(imX, "Forming initial guess using Median stack" )
            if not alive: return None
            xn = np.median(stack,axis=0)
            xn = xn.reshape(shape)
        elif methodInitialGuess==self.INITIAL_KALMAN:
            alive=callback(imX, "Forming initial guess using Kalman filter" )
            if not alive: return None

            xhat = np.median(stack, axis=0)   #A posteriori estimate of x
            P         = np.ones(xhat.shape)   #A posteriori error estimate
            xhatMinus = np.zeros(xhat.shape)  #A priori estimate of x
            Pminus    = np.zeros(xhat.shape)  #A priori error estimate
            K         = np.zeros(xhat.shape)  #Gain or blending factor
            error=np.std(stack,axis=0)
            Q = np.mean(error) #Process noise
            R = np.var(error)  #Estimate of measurement error variance
            for frame in stack:
                xhatMinus[:] = xhat
                Pminus[:] = P + Q
                K[:] = Pminus / (Pminus + R)
                xhat[:] = xhatMinus + K*(frame-xhatMinus)
                P[:] = (1. - K) * Pminus
            xn = xhat
            xn = xn.reshape(shape)
        
        #Estimate the initial noise as y = p*x + n => n = y - p*x
        print "Noise estimate:",
        noiseMean = np.zeros(shape)
        noiseMean = np.fft.rfft2(noiseMean,s=shape)
        noiseSigma=0
        for i in xrange(noOfFrames):
            yt = np.reshape(stack[i],shape)
            ytn = yt / np.mean(yt) * np.mean(xn)
            noise = ytn - xn
            noiseMean = noiseMean + np.abs(np.fft.rfft2(noise,s=shape))
            noiseSigma = noiseSigma + np.sqrt(np.mean(noise**2))
        noiseSpectra = noiseMean / noOfFrames
        noiseSpectra = noiseSpectra / np.max(noiseSpectra)
        noiseSigma = noiseSigma / noOfFrames
        NSR = noiseSigma**2/np.mean(xn**2) 
        print "Noise to Signal Ratio=%.1f%%" % (NSR*100.)
        imX = Image.fromarray(xn.astype(np.int8),"L")
        alive=callback(imX, "Noise to Signal Ratio=%.1f%%" % (NSR*100.) )
        if not alive: return None      
        
        if methodInitialGuess!=self.INITIAL_FIRSTFRAME:
            #Calculate initial guess at PSF using a gaussian of specified radius
            gaussParams = ( 1.,psfsize/2., psfsize/2.,  blurRadius/4., blurRadius/4., 0. )
            gaussPSF = self.twodgaussian(gaussParams,psfshape)
            gaussPSF = gaussPSF / np.sum(gaussPSF)            

            alive=callback(imX, "Deconvolving with Gaussian PSF of radius %.1f" % blurRadius )
            if not alive: return None
            #xn = self.LRDeconvolution(xn,gaussPSF,tolerance,iterations=Miterations,callback=callback,LRThreshold=LRThreshold,lambdaTV=lambdaTV, noiseSpectra=noiseSpectra)
            xn = self.TVDDeconvolution(xn,gaussPSF,alpha=noiseSigma)
            
        xn0 = np.array(xn)
        imX = Image.fromarray(xn0.astype(np.int8),"L")
        if debugPath!=None:
            imX.save("" + debugPath + "src/test/x0.png")
        
        #Save the input tile
        for i,ytflat in enumerate(sstack):  
            yt = np.reshape(ytflat,shape)
            sampleyt = yt[x0:x1,y0:y1]
            imSY = Image.fromarray(sampleyt.astype(np.int8),"L")
            if debugPath!=None:
                imSY.save("" + debugPath + "src/test/sy%03d.png" % i)
                              
        #PSF Fitting residual
        entropyxn0 = -np.mean(xn0[xn0>0.]*np.log(xn0[xn0>0.]))
        sharpnessxn0 = np.mean(np.array(np.gradient(xn0))**2)

        def residual(params):
            radius = blurRadius
            centre = params[0:2]
            coeffsReal = params[2:2+noOfShapeletCoeffs]
            coeffsImag = params[2+noOfShapeletCoeffs:]
            psfRes = self.shapelet(radius,centre,coeffsReal+1.j*coeffsImag,psfshape)
            psfSum = np.sum(psfRes)
            if psfSum!=0: 
                psfRes = psfRes / psfSum
            else:
                psfRes = np.zeros(psfshape)
                psfRes[psfsize/2,psfsize/2] = 1.
            
            #Score the deconvolution by comparing to a convolution back to the sample image
            paddedconv = filters.convolve(paddedx,psfRes)
            sampleconv = paddedconv[psfsize:-psfsize,psfsize:-psfsize]
            resid = np.sum( (sampleyt - sampleconv)**2 )/np.sum(sampleyt)
            
            if 1:
                #Bias the PSF choice to give high sharpness after deconvolution
                if 0:
                    paddedDeconv = self.LRDeconvolution(paddedyt, psfRes,iterations=3)
                else:
                    paddedDeconv = self.TVDDeconvolution(paddedyt, psfRes,alpha=noiseSigma/noOfFrames)
                
                #sharpness = np.mean(np.array(np.gradient(paddedDeconv))**2)
                #resid = resid - sharpness/sharpnessxn0 
                entropy = -np.mean(paddedDeconv[paddedDeconv>0.]*np.log(paddedDeconv[paddedDeconv>0.]))
                resid = resid - 10.*entropy/entropyxn0
            return resid

        #Form initial guess for the PSF
        centre = [psfsize/2,psfsize/2]
        noOfShapeletCoeffs = shapeletOrder*(shapeletOrder+1)/2
        psfparams = np.zeros( (noOfFrames, 2+2*noOfShapeletCoeffs) )
        psfparams[:,0] = psfsize/2
        psfparams[:,1] = psfsize/2
        psfparams[:,2:2+noOfShapeletCoeffs] =     -1 + 2.*np.random.random( (noOfFrames, noOfShapeletCoeffs) )
        psfparams[:,2+noOfShapeletCoeffs: ] = -np.pi + 2.*np.pi*np.random.random( (noOfFrames, noOfShapeletCoeffs) )
        
        for EMiter in xrange(EMiterations):
            #Retrieve the current solution tile
            samplex = xn[x0:x1,y0:y1]
            paddedx = xn[px0:px1,py0:py1]
            samplex0 = xn0[x0:x1,y0:y1]
            paddedx0 = xn0[px0:px1,py0:py1]

            psf = np.zeros((noOfFrames,psfsize,psfsize))
            for i in xrange(noOfFrames):
                #Extract the current frame from the sample stack
                yt = np.reshape(sstack[i],shape)
                sampleyt = yt[x0:x1 ,y0:y1]
                paddedyt = yt[px0:px1, py0:py1]
                psfp = psfparams[i]
                                    
                #Fit the PSF for the current frame
                rmax = psfsize/3.
                boundsRadius = [[1., psfsize]]
                boundsCentre = [[psfsize/2-rmax,psfsize/2+rmax], [psfsize/2-rmax,psfsize/2+rmax]]
                boundsCoeffsReal = np.array( ([-1.,1.])*noOfShapeletCoeffs )
                boundsCoeffsImag = np.array( ([-np.pi,np.pi])*noOfShapeletCoeffs )
                bounds = boundsCentre
                bounds = np.append(bounds, boundsCoeffsReal)
                bounds = np.append(bounds, boundsCoeffsImag)
                bounds = np.reshape(bounds, (2+2.*noOfShapeletCoeffs, 2))
                if 0:
                    psfp, val,d = fmin_l_bfgs_b(residual,psfp,approx_grad=True,bounds=bounds,epsilon=0.01,maxfun=Eiterations)
                else:
                    psfp = fmin_powell(residual,psfp,maxiter=Eiterations,ftol=0.01)
                    
                centre = psfp[0:2]
                coeffsReal = psfp[2:2+noOfShapeletCoeffs]
                coeffsImag = psfp[2+noOfShapeletCoeffs:]
                psf[i] = self.shapelet(blurRadius,centre,coeffsReal+1.j*coeffsImag,psfshape)
                    
                #Normalise the PSF
                psf[i] = psf[i]/np.sum(psf[i])
                                                                
                #Display progress
                if callback!=None:
                    psfStretched = psf[i] / np.max(psf[i])
                    xnH = np.array(yt)
                    xnH[psfx0:psfx1,psfy0:psfy1] = 255.*psfStretched
                    xnH = np.clip(xnH,0,255)
                    imX = Image.fromarray(xnH.astype(np.int8),"L")
                    alive=callback(imX, "Iteration %d. PSF for frame %d of %d" % (EMiter+1, i+1, noOfFrames) )
                    if not alive: return None

            #Clip and save the PSFs to file
            for i in xrange(noOfFrames):
                psfStretched = psf[i] / np.max(psf[i])
                allPSF[i,psfx0:psfx1,psfy0:psfy1] =  psfStretched
                imPSF = Image.fromarray((255*psfStretched).astype(np.int8),"L")
                if debugPath!=None:
                    imPSF.save("" + debugPath + "src/test/psf%03d.png" % i)

                #Save image of all of the PSFs to file
                imPSF = Image.fromarray((255*allPSF[i]).astype(np.int8),"L")
                if debugPath!=None:
                    imPSF.save("" + debugPath + "src/test/allpsf%03d.png" % i)
                    
            if 0:
                #Lucy-Richardson deconvolution 'M-step'
                paddedu = self.MFLR(Miterations,noOfFrames,sstack,shape,psf, 
                                    tolerance=tolerance,
                                    lambdaTV=lambdaTV,
                                    noiseSpectra=noiseSpectra,
                                    LRThreshold=LRThreshold, 
                                    waveletFilter=waveletFilter, 
                                    noiseSigma=noiseSigma,
                                    callback=callback)
                if paddedu==None: return None

            else:
                paddedu = self.MFTVD(sstack,shape,psf,alpha=noiseSigma/noOfFrames)

            sampleu = paddedu[psfsize:-psfsize,psfsize:-psfsize]
            samplex = xn[psfsize:-psfsize,psfsize:-psfsize]
            deltax = (sampleu - samplex)
            xn[psfsize:-psfsize,psfsize:-psfsize]=samplex + deltax

            #Estimate the noise level
            print "Noise estimate:",
            noiseMean = np.zeros(shape)
            noiseMean = np.fft.rfft2(noiseMean,s=shape)
            noiseSigma=0
            for i in xrange(noOfFrames):
                yt = np.reshape(stack[i],shape)
                ytn = yt / np.mean(yt) * np.mean(xn)
                noise = ytn - xn
                noiseMean = noiseMean + np.abs(np.fft.rfft2(noise,s=shape))
                noiseSigma = noiseSigma + np.sqrt(np.mean(noise**2))
            noiseSpectra = noiseMean / noOfFrames
            noiseSpectra = noiseSpectra / np.max(noiseSpectra)
            noiseSigma = noiseSigma / noOfFrames
            
            NSR = noiseSigma**2/np.mean(xn**2) 
            print "Noise to Signal Ratio=%.1f%%" % (NSR*100.)

            globalResidual = np.sqrt(np.max(deltax**2))/np.max(xn)
            print "Global residual=", globalResidual
            if globalResidual<tolerance: break
            
            #Store the result image
            xClipped = np.clip(xn, 0,255)
            imX = Image.fromarray(xClipped.astype(np.int8),"L")
            if debugPath!=None:
                imX.save("" + debugPath + "src/test/x.png")
                        
            if callback!=None:
                alive=callback(imX, "Iteration %d.  Max change=%.1f%%.  NSR=%.2f%%" % (EMiter+1,globalResidual*100.,NSR*100.) )
                if not alive: return None
            
        xnExtra=[]
        if stackExtra!=None:
            for i,s in enumerate(stackExtra):    
                if 0:
                    paddedu = self.MFLR(Miterations,noOfFrames,s,shape,psf, tolerance=tolerance,LRThreshold=LRThreshold, waveletFilter=waveletFilter, noiseSigma=noiseSigma,callback=callback)
                else:
                    paddedu = self.MFTVD(sstack,shape,psf,alpha=noiseSigma/noOfFrames)
                xnExtra.append(paddedu)
            xnExtra.append(xn)
            xnExtra=np.array(xnExtra)
        else:
            xnExtra=xn
        
        if callback!=None:
            alive=callback(imX, "MFBD Complete. Result saved to x.png" )
            if not alive: return None

        return xnExtra

    def TVDDeconvolution(self,img,kernel,alpha=0.01, crop=True, fftshift=True, fft_pad=True, psf_pad=False):
        """
        Deconvolve an image with a kernel.  Returns something the size of an image.
        Assumes image & kernel are centred
        *NOTE* Order matters; the kernel should be second.
        """

        imgshape = img.shape
        kernshape = kernel.shape
        # find ideal size (power of 2) for fft.  Can add shapes because they are tuples
        if fft_pad:
            if psf_pad: 
                # add the X dimensions and Y dimensions and then take the max (bigger)
                fsize = 2**np.ceil(np.log2(np.max(np.array(imgshape)+np.array(kernshape)))) 
            else: 
                # add the shape lists (max of a list of length 4) (smaller)
                fsize = 2**np.ceil(np.log2(np.max(imgshape+kernshape)))
            newshape = np.array([fsize,fsize])
        else:
            if psf_pad:
                newshape = np.array(imgshape)+np.array(kernshape) # just add the biggest dimensions
            else:
                newshape = np.array([np.max([imgshape[0],kernshape[0]]),np.max([imgshape[1],kernshape[1]])]) 

        centerx, centery = newshape/2.
        imgquarter1x, imgquarter1y = centerx - imgshape[0]/2.,centery - imgshape[1]/2.
        imgquarter3x, imgquarter3y = imgquarter1x + imgshape[0],imgquarter1y + imgshape[1]
        kernelquarter1x, kernelquarter1y = centerx - kernshape[0]/2.,centery - kernshape[1]/2.
        kernelquarter3x, kernelquarter3y = kernelquarter1x + kernshape[0],kernelquarter1y + kernshape[1]
        bigimg = np.zeros(newshape,dtype=np.float32)
        bigkernel = np.zeros(newshape,dtype=np.float32)
        bigimg[imgquarter1x:imgquarter3x,imgquarter1y:imgquarter3y] = img
        bigkernel[kernelquarter1x:kernelquarter3x,kernelquarter1y:kernelquarter3y] = kernel 

        yhat = rfft2(bigimg)
        phat = rfft2(bigkernel)
        uhat = yhat * np.conj(phat) / (phat*np.conj(phat) + alpha**2)

        rifft = np.fft.fftshift( irfft2( uhat ) ) 
        rifft = rifft / np.sum(rifft) * np.sum(img)

        if crop:
            result = rifft[ imgquarter1x:imgquarter3x, imgquarter1y:imgquarter3y ]
            return result
        else:
            return rifft        


    def MFTVD(self,stack,shape,psfstack,alpha=0.01):
        noOfFrames = stack.shape[0]
        ustack = np.zeros( (noOfFrames, shape[0],shape[1]) )
        for i in xrange(noOfFrames): 
            #Extract the current frame
            yt = np.reshape(stack[i],shape) 

            #Deconvolve
            ustack[i] = self.TVDDeconvolution(yt,psfstack[i],alpha)
        
        u = np.median(ustack,axis=0)
        return u

    #Multi-frame Lucy-Richardson
    def MFLR(self,Miterations,noOfFrames,stack,shape,psfstack,initialGuess=None,tolerance=1./255.,LRThreshold=0, lambdaTV=0, waveletFilter=False, noiseSigma=0.,noiseSpectra=None,callback=None):
        print "Lucy-Richardson M-step"
        eps = 1.e-8
        psfsize=psfstack.shape[1]
            
        if initialGuess==None:
            u = np.mean(stack[0])*np.ones(shape)
        else:
            u = initialGuess
            
        for Miter in range(Miterations):
            u0 = np.array(u)
            ustack = np.zeros( (noOfFrames, u.shape[0],u.shape[1]) )
            for i in xrange(noOfFrames): 
                #Extract the current frame
                yt = np.reshape(stack[i],shape) 

                #Lucy-richardson iteration
                c = filters.convolve(u,psfstack[i])
                c[c==0] = eps
                unew = u*filters.correlate(yt/c,psfstack[i])

                #Spectral filter
                if noiseSpectra!=None:
                    f = np.fft.rfft2(unew,s=unew.shape)
                    c = (1.-lambdaTV*noiseSpectra)
                    msk = c>0
                    f[msk] = f[msk]/c[msk]
                    unew = np.fft.irfft2(f,s=unew.shape)
                    unew = unew/np.sum(unew) * np.sum(u)

                #Total-variation scaling
                elif Miter>0 and lambdaTV!=0:
                    gradu = np.array(np.gradient(u))
                    gradu = gradu/np.mean(gradu**2)**0.5
                    divgradu = np.gradient(gradu[0])[0] + np.gradient(gradu[1])[1]
                    filterTV =  1./(1. - lambdaTV*divgradu)
                    unew = np.convolve(unew, filterTV)
                    
                #Change thresholding                
                if Miter>0 and LRThreshold>0:
                    mask = np.abs(unew-u) > LRThreshold
                    unew[mask] = u[mask]
                
                #Store the deconvoled frame in the stack
                ustack[i] = unew
                
            #Find the median of the all frames in the stack at the 
            #  current iteration of the Lucy-Richardson
            u = np.median(ustack,axis=0)

            #Apply a wavelet noise filter
            if waveletFilter and noiseSigma>0:
                s = np.log(Miter+1)*noiseSigma/(noOfFrames*np.log(Miterations))
                u = filters.waveletDenoise(u,s)
                
            if callback!=None:
                uClipped = np.clip(u,0,255)
                xnH = np.array(yt)
                xnH[psfsize:-psfsize,psfsize:-psfsize]=uClipped[psfsize:-psfsize,psfsize:-psfsize]
                imX = Image.fromarray(xnH.astype(np.int8),"L")                        
                alive=callback(imX, "Deconvolution step %d of %d" % (Miter+1,Miterations) )
                if not alive: return None
            
            Mresidual = np.sqrt(np.max((u-u0)**2))/np.max(u)
            print "M-step residual (%d)=%f (<%f?)" % (Miter, Mresidual, tolerance)
            if Mresidual<tolerance: break

        return u

    def LRDeconvolution(self,paddedyt,kern,tolerance=1./255.,iterations=10,callback=None,LRThreshold=None,lambdaTV=0.,waveletFilter=False,noiseSigma=0.,noiseSpectra=None): 
        #Lucy-Richardson E-step
        eps = 1.e-8
        iterations=int(iterations)
        u = np.ones(paddedyt.shape)*np.mean(paddedyt)
        for Miter in range(iterations):
            c = filters.convolve(u,kern)
            c[c==0] = eps
            unew = u*filters.correlate(paddedyt/c,kern)
            
            #Apply a wavelet noise filter
            if waveletFilter and noiseSigma>0:
                unew = filters.waveletDenoise(unew,noiseSigma)
                
            #Spectral filter
            if noiseSpectra!=None:
                f = np.fft.rfft2(unew,s=unew.shape)
                print lambdaTV,np.max(noiseSpectra)
                c = (1.-lambdaTV*noiseSpectra)
                msk = c>0
                f[msk] = f[msk]/c[msk]
                unew = np.fft.irfft2(f,s=unew.shape)
                unew = unew/np.sum(unew) * np.sum(u)
                
            #Total-variation filtering
            elif Miter>0 and lambdaTV!=0:
                gradu = np.array(np.gradient(u))
                gradu = gradu/np.mean(gradu**2)**0.5
                divgradu = np.gradient(gradu[0])[0] + np.gradient(gradu[1])[1]
                unew =  unew/(1. - lambdaTV*divgradu)

            #Change thresholding                
            if Miter>0 and LRThreshold>0:
                mask = np.abs(unew-u) > LRThreshold
                unew[mask] = u[mask]
                
            du = unew-u
            u = u + du
            Mresidual = np.sqrt(np.max(du**2))/np.max(u)
            
            if callback!=None:
                uClipped = np.clip(u,0,255)
                imX = Image.fromarray(uClipped.astype(np.int8),"L")                        
                alive=callback(imX, "LRiteration %d" % (Miter+1) )
                if not alive: return u

            if Mresidual<tolerance: break                                    
            
        return u
              
    def twodgaussian(self,inpars, shape):
        """Returns a 2d image of a gaussian functionp.
            inpars = (brightness,center_x,center_y,width_x,width_y,rota)
            shape - a 2-parameter list
            """
        center_y, center_x = inpars[1],inpars[2]
        center_x = float(center_x)
        center_y = float(center_y)
        width_x, width_y = inpars[3],inpars[4]
        width_x = max(0.1,abs(float(width_x)))
        width_y = max(0.1,abs(float(width_y)))
        rota = inpars[5]
        rota = float(rota)
        rcen_x = center_x * np.cos(rota) - center_y * np.sin(rota)
        rcen_y = center_x * np.sin(rota) + center_y * np.cos(rota)
        brightness = inpars[0]
        
        def rotgauss(x,y):
            xp = x * np.cos(rota) - y * np.sin(rota)
            yp = x * np.sin(rota) + y * np.cos(rota)
            g = np.exp( -(((rcen_x-xp)/width_x)**2 + ((rcen_y-yp)/width_y)**2)/2.)
            return g
        
        g = brightness*rotgauss(*np.indices(shape))
        if np.sum(g)<1.e-9:
            print "Zero gaussian:", inpars 
            g=np.zeros(shape)
        return g
        
    def shapelet(self,radius,centre,coeffs,shape,aspect=1.):
        shplt = np.zeros(shape)
        scale = [1./aspect,aspect]
        xx,yy = np.meshgrid(np.arange(shape[0]),np.arange(shape[1]))
        xx,yy = scale[0]*xx-centre[0], scale[1]*yy-centre[1]
        rr    = np.sqrt(xx**2 + yy**2)
        tt    = np.arctan2(yy,xx)

        #The optimal shapelet scale given by Richard Massey is around half the image radius
        beta = radius/2.

        n0=1
        m0=-1
        for i,cr in enumerate(coeffs[:-1]):
            ci=coeffs[i+1]
            na=int((n0-np.abs(m0))/2)
            nb=int((n0+np.abs(m0))/2)
            if nb!=0:
                L=special.genlaguerre(n=na,alpha=np.abs(m0))
                ff=lambda x: float(factorial(x))
                norm=(-1.)**na / beta**(np.abs(m0)+1) * np.sqrt( ff(na) / ff(nb) / np.pi ) 
                exp0=lambda r,th: norm * r**np.abs(m0) * L(r**2./beta**2.) * np.exp(-r**2./(2.*beta**2.)) * np.exp(-1j*m0*th)
                shplt=shplt + np.real((cr+ci*1.j)*exp0(rr,tt))

            m0 = m0 + 2
            if m0>n0:
                n0 = n0 + 1
                m0 = -n0
        
        shplt[shplt<0] = 0.
        return shplt
                    
#If executed as main, then generate a test case in the current directory
if __name__ == "__main__":
    import sys,os
    
    if len(sys.argv)!=2:
        sys.exit('Usage: %s image-file' % sys.argv[0])

    if not os.path.exists(sys.argv[1]):
        sys.exit('ERROR: Image %s was not found!' % sys.argv[1])
        
    psfsize=16
    psfshape=(psfsize,psfsize)
    mfbd = MFBD()

    basename=os.path.splitext(sys.argv[1])[0]
    im = Image.open(sys.argv[1])
    im = ImageOps.expand(im, psfsize)

    nim = np.asarray(im)
    if len(nim.shape)==3:
        nim = np.mean(nim, axis=2)
    greyIm = Image.fromarray(nim.astype(np.int8),"L")
    greyIm.save("%s-grey.png" % basename)

    NSR = 0.1
    n=20
    shiftx,shifty=6.,6.
    sx,sy=4.,1.
    for i in range(n):
        psf = np.zeros(psfsize)
        for psfNo in range(2):        
            theta = np.random.random() * np.pi
            ox = -sx + float(i)/float(n-1)*2.*sx
            oy = -sy + float(i)/float(n-1)*2.*sy
            
            psfparams=[1.0,psfsize/2+ox,psfsize/2+oy,sx,sy,theta]
            psf = psf + mfbd.twodgaussian(psfparams,psfshape)

        psf = psf/np.sum(psf)
        psfIm = Image.fromarray(psf.astype(np.int8),"L")
        
        noiseLevel=np.max(nim)*NSR
        noise = (np.random.random(nim.shape)-0.5)*noiseLevel
        
        cnim = mfbd.convolve(nim,psf) + noise
        cnim[cnim<0] = 0
        cnim[cnim>255] = 255
        cim = Image.fromarray(cnim.astype(np.int8),"L")
        cim.save("%s-%.3f.png" % (basename, 180./np.pi * theta))

